-module(optimizer).

-export([parse_transform/2]).

parse_transform(Tree, _Options) ->
	FunctionFactories = [{Name, lists:map(fun({_, _, VarName}) -> VarName end, Vars), Clauses} || {function, _, Name, _, [{clause, _, Vars, [], [{'fun', _, {clauses, Clauses}}]}]} <- Tree],

	Transform = inject_factories(Tree, FunctionFactories),

	Transform.

inject_factories(AST, Factories) ->
	walk_ast(fun
		({function, Line, Name, Arity, [{clause, _, _, _, [{call, _, {'fun', _, {clauses, Clauses}}, _}]}]}) ->
			{function, Line, Name, Arity, Clauses};
		({call, Line, {atom, _, FactoryName}, Arguments} = Call) ->
			case keyfind(FactoryName, 1, Factories) of
				{FactoryName, Vars, Clauses0} ->
					Clauses = inject_values(Clauses0, lists:zip(Vars, Arguments)),
					{'fun', Line, {clauses, Clauses}};
				_ -> Call
			end;
		(ASTItem) -> ASTItem
	end, AST).

inject_values(AST, Replacements) ->
	walk_ast(fun
		({var, _, Variable} = Item) ->
			case keyfind(Variable, 1, Replacements) of
				false -> Item;
				{_, Replacement} -> Replacement
			end;
		(ASTItem) -> ASTItem
	end, AST).

walk_ast(F, ASTItem) when is_tuple(ASTItem) ->
	hd(walk_ast(F, [ASTItem]));
walk_ast(F, AST) ->
	lists:map(fun
		(List) when is_list(List) -> [walk_ast(F, ListItem) || ListItem <- List];
		({function, Line, Name, Arity, Clauses}) -> F({function, Line, Name, Arity, walk_ast(F, Clauses)});
		({'fun', Line, {clauses, Clauses}}) -> F({'fun', Line, {clauses, walk_ast(F, Clauses)}});
		({'if', Line, Clauses}) -> F({'if', Line, walk_ast(F, Clauses)});
		({'case', Line, Test, Clauses}) -> F({'case', Line, walk_ast(F, Test), walk_ast(F, Clauses)});
		({'receive', Line, Clauses}) -> F({'receive', Line, walk_ast(F, Clauses)});
		({'receive', Line, Clauses1, TimeAfter, Clauses2}) -> F({'receive', Line, walk_ast(F, Clauses1), TimeAfter, walk_ast(F, Clauses2)});
		({clause, Line, Pattern, Guards, Expressions}) -> F({clause, Line, walk_ast(F, Pattern), walk_ast(F, Guards), walk_ast(F, Expressions)});
		({record, Line, Record, Fields}) -> F({record, Line, Record, walk_ast(F, Fields)});
		({record_field, Line, Field, Contents}) -> F({record_field, Line, walk_ast(F, Field), walk_ast(F, Contents)});
		({record_field, Line, Variable, Record, Field}) -> F({record_field, Line, Variable, Record, walk_ast(F, Field)});

		({call, Line, Fun, Arguments}) -> F({call, Line, walk_ast(F, Fun), walk_ast(F, Arguments)});
		({op, Line, BinaryOp, Left, Right}) -> F({op, Line, BinaryOp, walk_ast(F, Left), walk_ast(F, Right)});
		({match, Line, Left, Right}) -> F({match, Line, walk_ast(F, Left), walk_ast(F, Right)});

		({tuple, Line, Items}) -> F({tuple, Line, walk_ast(F, Items)});
		({cons, Line, Head, Tail}) -> F({cons, Line, walk_ast(F, Head), walk_ast(F, Tail)});
		({bin, Line, Elements}) -> F({bin, Line, walk_ast(F, Elements)});
		({bin_element, Line, Name, Size, Type}) -> F({bin_element, Line, walk_ast(F, Name), Size, Type});

		({attribute, Line, Attribute, Data}) -> F({attribute, Line, Attribute, Data});
		({var, Line, Atom}) -> F({var, Line, Atom});
		({atom, Line, Atom}) -> F({atom, Line, Atom});
		({integer, Line, Integer}) -> F({integer, Line, Integer});
		({remote, Line, Module, Function}) -> F({remote, Line, Module, Function});
		({char, Line, Char}) -> F({char, Line, Char});
		({nil, Line}) -> F({nil, Line});
		({'fun', Line, {function, Function, Arity}}) -> F({'fun', Line, {function, Function, Arity}});
		({eof, Line}) -> F({eof, Line});
		
		(ASTItem) -> io:format("unkown AST: ~p~n", [ASTItem]), F(ASTItem)
	end, AST).

keyfind(Key, Index, TupleList) ->
	lists:foldl(fun
		(Tuple, false) when element(Index, Tuple) =:= Key -> Tuple;
		(_, Result) -> Result
	end, false, TupleList).

